"""
牛顿法数值分析

问题描述：
    牛顿法（Newton's Method）是一种求解非线性方程根的数值方法，
    也可用于函数优化问题。它基于函数的泰勒展开和线性近似的思想，
    通过迭代逼近方程的根或函数的极值点。

算法原理：
    牛顿法的基本思想：
    1. 求根问题：f(x) = 0
       迭代公式：x_{n+1} = x_n - f(x_n)/f'(x_n)
    2. 优化问题：min f(x)
       迭代公式：x_{n+1} = x_n - f'(x_n)/f''(x_n)
    3. 收敛条件：|x_{n+1} - x_n| < ε 或 |f(x_n)| < ε

核心概念：
    - 导数计算：解析导数或数值导数
    - 收敛性：二次收敛（接近根时）
    - 初值选择：影响收敛性和收敛到哪个根
    - 收敛判据：绝对误差、相对误差

应用场景：
    牛顿法广泛应用于：
    - 非线性方程求解
    - 函数优化
    - 平方根计算
    - 参数估计
    - 工程计算

作者：斯黄
日期：2025年3月4日
"""

import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import fsolve, minimize_scalar
import sympy as sp

class NewtonMethod:
    def __init__(self, tolerance=1e-8, max_iterations=100):
        self.tolerance = tolerance
        self.max_iterations = max_iterations
    
    def newton_root(self, f, df, x0, verbose=False):
        """
        使用牛顿法求解方程 f(x) = 0 的根
        
        参数:
        f: 目标函数
        df: 目标函数的导数
        x0: 初始猜测值
        verbose: 是否输出迭代过程
        """
        x = x0
        iterations = []
        
        for i in range(self.max_iterations):
            fx = f(x)
            dfx = df(x)
            
            if abs(dfx) < 1e-12:  # 避免除零
                if verbose:
                    print(f"导数为零，停止迭代")
                break
            
            x_new = x - fx / dfx
            iterations.append((i, x, fx, x_new))
            
            if verbose:
                print(f"迭代 {i+1}: x = {x:.8f}, f(x) = {fx:.8f}, x_new = {x_new:.8f}")
            
            # 检查收敛性
            if abs(x_new - x) < self.tolerance or abs(fx) < self.tolerance:
                if verbose:
                    print(f"收敛！找到根: x = {x_new:.8f}")
                return x_new, i+1, iterations
            
            x = x_new
        
        if verbose:
            print(f"达到最大迭代次数，可能未收敛")
        return x, self.max_iterations, iterations
    
    def newton_optimization(self, f, df, ddf, x0, verbose=False):
        """
        使用牛顿法求解优化问题 min f(x)
        
        参数:
        f: 目标函数
        df: 一阶导数
        ddf: 二阶导数（Hessian）
        x0: 初始点
        """
        x = x0
        iterations = []
        
        for i in range(self.max_iterations):
            dfx = df(x)
            ddfx = ddf(x)
            
            if abs(ddfx) < 1e-12:  # 避免除零
                if verbose:
                    print(f"二阶导数为零，停止迭代")
                break
            
            x_new = x - dfx / ddfx
            iterations.append((i, x, f(x), dfx, x_new))
            
            if verbose:
                print(f"迭代 {i+1}: x = {x:.8f}, f(x) = {f(x):.8f}, f'(x) = {dfx:.8f}")
            
            # 检查收敛性
            if abs(x_new - x) < self.tolerance or abs(dfx) < self.tolerance:
                if verbose:
                    print(f"收敛！找到极值点: x = {x_new:.8f}, f(x) = {f(x_new):.8f}")
                return x_new, i+1, iterations
            
            x = x_new
        
        if verbose:
            print(f"达到最大迭代次数，可能未收敛")
        return x, self.max_iterations, iterations
    
    def numerical_derivative(self, f, x, h=1e-5):
        """数值计算导数"""
        return (f(x + h) - f(x - h)) / (2 * h)
    
    def numerical_second_derivative(self, f, x, h=1e-5):
        """数值计算二阶导数"""
        return (f(x + h) - 2*f(x) + f(x - h)) / (h**2)

# 示例函数定义
def example_functions():
    """定义示例函数及其导数"""
    
    # 函数1: f(x) = x^3 - 2x - 5
    def f1(x):
        return x**3 - 2*x - 5
    
    def df1(x):
        return 3*x**2 - 2
    
    def ddf1(x):
        return 6*x
    
    # 函数2: f(x) = e^x - 3x
    def f2(x):
        return np.exp(x) - 3*x
    
    def df2(x):
        return np.exp(x) - 3
    
    def ddf2(x):
        return np.exp(x)
    
    # 函数3: f(x) = x^4 - 4x^3 + 6x^2 - 4x + 1 (优化问题)
    def f3(x):
        return x**4 - 4*x**3 + 6*x**2 - 4*x + 1
    
    def df3(x):
        return 4*x**3 - 12*x**2 + 12*x - 4
    
    def ddf3(x):
        return 12*x**2 - 24*x + 12
    
    return (f1, df1, ddf1), (f2, df2, ddf2), (f3, df3, ddf3)

def demonstrate_root_finding():
    """演示牛顿法求根"""
    print("=== 牛顿法求根演示 ===")
    
    newton = NewtonMethod()
    (f1, df1, _), (f2, df2, _), _ = example_functions()
    
    # 示例1: x^3 - 2x - 5 = 0
    print("\n1. 求解方程 x³ - 2x - 5 = 0")
    print("不同初值的收敛情况:")
    
    initial_values = [1.0, 2.0, 3.0, -1.0]
    for x0 in initial_values:
        root, iterations, _ = newton.newton_root(f1, df1, x0)
        print(f"初值 x0 = {x0:>4.1f}: 根 = {root:.6f}, 迭代次数 = {iterations}")
        print(f"             验证 f({root:.6f}) = {f1(root):.2e}")
    
    # 示例2: e^x - 3x = 0
    print("\n2. 求解方程 e^x - 3x = 0")
    x0 = 1.0
    root, iterations, iter_history = newton.newton_root(f2, df2, x0, verbose=True)
    
    return iter_history

def demonstrate_optimization():
    """演示牛顿法优化"""
    print("\n=== 牛顿法优化演示 ===")
    
    newton = NewtonMethod()
    _, _, (f3, df3, ddf3) = example_functions()
    
    print("求解 f(x) = x⁴ - 4x³ + 6x² - 4x + 1 的极值")
    
    initial_values = [0.0, 1.5, 3.0]
    for x0 in initial_values:
        extremum, iterations, _ = newton.newton_optimization(f3, df3, ddf3, x0)
        print(f"初值 x0 = {x0:>4.1f}: 极值点 = {extremum:.6f}, f(x) = {f3(extremum):.6f}, 迭代次数 = {iterations}")

def square_root_newton():
    """使用牛顿法计算平方根"""
    print("\n=== 牛顿法计算平方根 ===")
    
    def sqrt_newton(a, x0=1.0, tolerance=1e-10):
        """使用牛顿法计算√a"""
        x = x0
        iterations = 0
        
        while True:
            x_new = 0.5 * (x + a / x)
            iterations += 1
            
            if abs(x_new - x) < tolerance:
                break
            x = x_new
            
            if iterations > 50:  # 防止无限循环
                break
        
        return x_new, iterations
    
    test_numbers = [2, 7, 10, 25, 100]
    
    for a in test_numbers:
        result, iters = sqrt_newton(a)
        actual = np.sqrt(a)
        error = abs(result - actual)
        
        print(f"√{a:>3}: 牛顿法 = {result:.10f}, 实际值 = {actual:.10f}")
        print(f"      误差 = {error:.2e}, 迭代次数 = {iters}")

def visualize_newton_method():
    """可视化牛顿法收敛过程"""
    print("\n=== 牛顿法可视化 ===")
    
    # 函数: f(x) = x^3 - 2x - 5
    def f(x):
        return x**3 - 2*x - 5
    
    def df(x):
        return 3*x**2 - 2
    
    newton = NewtonMethod()
    x0 = 2.0
    root, iterations, iter_history = newton.newton_root(f, df, x0)
    
    # 绘制函数和牛顿法迭代过程
    x = np.linspace(1.5, 3.0, 1000)
    y = f(x)
    
    plt.figure(figsize=(14, 10))
    
    # 主图：函数曲线和切线
    plt.subplot(2, 2, 1)
    plt.plot(x, y, 'b-', linewidth=2, label='f(x) = x³ - 2x - 5')
    plt.axhline(y=0, color='k', linestyle='--', alpha=0.3)
    plt.axvline(x=root, color='r', linestyle='--', alpha=0.7, label=f'根: x = {root:.4f}')
    
    # 绘制牛顿法迭代过程
    for i, (iter_num, xi, fxi, xi_new) in enumerate(iter_history[:5]):  # 只显示前5次迭代
        # 绘制点(xi, f(xi))
        plt.plot(xi, fxi, 'ro', markersize=6)
        
        # 绘制切线
        slope = df(xi)
        x_tangent = np.linspace(xi - 1, xi + 1, 100)
        y_tangent = fxi + slope * (x_tangent - xi)
        plt.plot(x_tangent, y_tangent, 'g--', alpha=0.7)
        
        # 绘制垂直线到x轴
        plt.plot([xi_new, xi_new], [0, f(xi_new)], 'r:', alpha=0.7)
        
        plt.annotate(f'x{i}', (xi, fxi), xytext=(5, 5), textcoords='offset points')
    
    plt.xlabel('x')
    plt.ylabel('f(x)')
    plt.title('牛顿法求根可视化')
    plt.legend()
    plt.grid(True, alpha=0.3)
    
    # 收敛过程
    plt.subplot(2, 2, 2)
    x_values = [x0] + [item[3] for item in iter_history]
    iterations_plot = range(len(x_values))
    plt.plot(iterations_plot, x_values, 'bo-', linewidth=2, markersize=6)
    plt.axhline(y=root, color='r', linestyle='--', label=f'真实根: {root:.6f}')
    plt.xlabel('迭代次数')
    plt.ylabel('x值')
    plt.title('收敛过程')
    plt.legend()
    plt.grid(True, alpha=0.3)
    
    # 误差分析
    plt.subplot(2, 2, 3)
    errors = [abs(x - root) for x in x_values]
    plt.semilogy(iterations_plot, errors, 'go-', linewidth=2, markersize=6)
    plt.xlabel('迭代次数')
    plt.ylabel('误差 (对数尺度)')
    plt.title('收敛误差')
    plt.grid(True, alpha=0.3)
    
    # 函数值变化
    plt.subplot(2, 2, 4)
    f_values = [f(x) for x in x_values]
    plt.semilogy(iterations_plot, [abs(fv) for fv in f_values], 'mo-', linewidth=2, markersize=6)
    plt.xlabel('迭代次数')
    plt.ylabel('|f(x)| (对数尺度)')
    plt.title('函数值收敛')
    plt.grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.show()

def compare_methods():
    """比较不同求根方法"""
    print("\n=== 求根方法比较 ===")
    
    def f(x):
        return x**3 - 2*x - 5
    
    def df(x):
        return 3*x**2 - 2
    
    # 牛顿法
    newton = NewtonMethod()
    x0 = 2.0
    newton_root, newton_iters, _ = newton.newton_root(f, df, x0)
    
    # scipy的求根方法
    scipy_root = fsolve(f, x0)[0]
    
    # 二分法实现
    def bisection_method(f, a, b, tolerance=1e-8):
        iterations = 0
        while (b - a) / 2 > tolerance:
            c = (a + b) / 2
            if f(c) == 0:
                break
            elif f(a) * f(c) < 0:
                b = c
            else:
                a = c
            iterations += 1
        return (a + b) / 2, iterations
    
    bisection_root, bisection_iters = bisection_method(f, 1.5, 3.0)
    
    print(f"牛顿法:     根 = {newton_root:.8f}, 迭代次数 = {newton_iters}")
    print(f"SciPy方法:  根 = {scipy_root:.8f}")
    print(f"二分法:     根 = {bisection_root:.8f}, 迭代次数 = {bisection_iters}")
    print(f"真实根验证: f({newton_root:.8f}) = {f(newton_root):.2e}")

def multivariate_newton():
    """多变量牛顿法示例"""
    print("\n=== 多变量牛顿法 ===")
    
    def system_equations(vars):
        """二元非线性方程组"""
        x, y = vars
        return np.array([
            x**2 + y**2 - 4,      # x² + y² = 4 (圆)
            x*y - 1               # xy = 1 (双曲线)
        ])
    
    def jacobian(vars):
        """雅可比矩阵"""
        x, y = vars
        return np.array([
            [2*x, 2*y],
            [y, x]
        ])
    
    def multivariate_newton_step(vars, max_iter=20, tolerance=1e-8):
        """多变量牛顿法迭代"""
        for i in range(max_iter):
            f_val = system_equations(vars)
            j_val = jacobian(vars)
            
            # 求解线性系统 J * delta = -f
            try:
                delta = np.linalg.solve(j_val, -f_val)
            except np.linalg.LinAlgError:
                print("雅可比矩阵奇异")
                break
            
            vars_new = vars + delta
            
            if np.linalg.norm(delta) < tolerance:
                print(f"收敛于第 {i+1} 次迭代")
                return vars_new, i+1
            
            vars = vars_new
        
        return vars, max_iter
    
    # 不同初值
    initial_guesses = [(1.5, 1.5), (-1.5, -1.5), (2.0, 0.5)]
    
    for x0, y0 in initial_guesses:
        print(f"\n初值: ({x0}, {y0})")
        solution, iters = multivariate_newton_step(np.array([x0, y0]))
        x_sol, y_sol = solution
        
        # 验证解
        residual = system_equations(solution)
        print(f"解: x = {x_sol:.6f}, y = {y_sol:.6f}")
        print(f"迭代次数: {iters}")
        print(f"残差: [{residual[0]:.2e}, {residual[1]:.2e}]")

if __name__ == "__main__":
    # 运行所有演示
    iter_history = demonstrate_root_finding()
    demonstrate_optimization()
    square_root_newton()
    compare_methods()
    multivariate_newton()
    
    # 可视化
    visualize_newton_method() 